{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pynvml import *\n",
    "\n",
    "nvmlInit()\n",
    "vram = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(1)).free/1024.**2\n",
    "print('GPU1 Memory: %dMB' % vram)\n",
    "if vram < 8000:\n",
    "    raise Exception('GPU Memory too low')\n",
    "nvmlShutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import h5py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import *\n",
    "from collections import Counter\n",
    "import seaborn as sns\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import re\n",
    "import time\n",
    "import random\n",
    "\n",
    "from keras.layers import *\n",
    "from keras.models import *\n",
    "from keras.optimizers import *\n",
    "from keras.regularizers import l2\n",
    "from keras.utils.vis_utils import model_to_dot\n",
    "import keras.backend as K\n",
    "from make_parallel import make_parallel\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "IMAGE_DIR = 'image_contest_level_2'\n",
    "\n",
    "MODEL_NAME = 'model_l2加层生成器2_0.997656.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('image_contest_level_2/labels.txt', sep=' ', header=None)\n",
    "characters = u'0123456789()+-*/=君不见黄河之水天上来奔流到海复回烟锁池塘柳深圳铁板烧; '\n",
    "\n",
    "labels_len = np.array(map(lambda x:len(x.decode('utf-8')), df[0]))\n",
    "n_len = 45\n",
    "n, width, height, n_class, channels = 100000, 900, 72, len(characters), 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def decode(out):\n",
    "    return ''.join([characters[x] for x in out if x < n_class-1 and x > -1])\n",
    "\n",
    "def disp2(img):\n",
    "    cv2.imwrite('a.png', img)\n",
    "    return Image('a.png')\n",
    "\n",
    "def disp(img, txt=None, first=False):\n",
    "    global index\n",
    "    if first:\n",
    "        index = 1\n",
    "        plt.figure(figsize=(16, 9))\n",
    "    else:\n",
    "        index += 1\n",
    "    plt.subplot(4, 1, index)\n",
    "    if len(img.shape) == 2:\n",
    "        plt.imshow(img, cmap='gray')\n",
    "    else:\n",
    "        plt.imshow(img[:,:,::-1])\n",
    "    if txt:\n",
    "        plt.title(txt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 读取测试集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.zeros((n, width, height, channels), dtype=np.uint8)\n",
    "\n",
    "for i in tqdm(range(n)):\n",
    "    img = cv2.imread('crop_split_test/%d.png'%i).transpose(1, 0, 2)\n",
    "    a, b, _ = img.shape\n",
    "    X[i, :a, :b] = img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 载入模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ctc_lambda_func(args):\n",
    "    y_pred, labels, input_length, label_length = args\n",
    "    y_pred = y_pred[:, 2:, :]\n",
    "    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)\n",
    "\n",
    "rnn_size = 128\n",
    "\n",
    "l2_rate = 1e-5\n",
    "\n",
    "input_tensor = Input((width, height, 3))\n",
    "x = input_tensor\n",
    "for i, n_cnn in enumerate([2, 3, 6]):\n",
    "    for j in range(n_cnn):\n",
    "        x = Conv2D(32*2**i, (3, 3), padding='same', kernel_initializer='he_normal', kernel_regularizer=l2(l2_rate))(x)\n",
    "        x = BatchNormalization()(x)\n",
    "        x = Activation('relu')(x)\n",
    "    x = MaxPooling2D((2, 2))(x)\n",
    "\n",
    "cnn_model = Model(input_tensor, x, name='cnn')\n",
    "\n",
    "input_tensor = Input((width, height, 3))\n",
    "x = cnn_model(input_tensor)\n",
    "\n",
    "conv_shape = x.get_shape().as_list()\n",
    "rnn_length = conv_shape[1]\n",
    "rnn_dimen = conv_shape[3]*conv_shape[2]\n",
    "\n",
    "print conv_shape, rnn_length, rnn_dimen\n",
    "\n",
    "x = Reshape(target_shape=(rnn_length, rnn_dimen))(x)\n",
    "rnn_length -= 2\n",
    "\n",
    "x = Dense(rnn_size, kernel_initializer='he_normal', kernel_regularizer=l2(l2_rate))(x)\n",
    "x = BatchNormalization()(x)\n",
    "x = Activation('relu')(x)\n",
    "x = Dropout(0.25)(x)\n",
    "\n",
    "gru_1 = GRU(rnn_size, implementation=2, return_sequences=True, name='gru1')(x)\n",
    "gru_1b = GRU(rnn_size, implementation=2, return_sequences=True, go_backwards=True, name='gru1_b')(x)\n",
    "gru1_merged = add([gru_1, gru_1b])\n",
    "\n",
    "gru_2 = GRU(rnn_size, implementation=2, return_sequences=True, name='gru2')(gru1_merged)\n",
    "gru_2b = GRU(rnn_size, implementation=2, return_sequences=True, go_backwards=True, name='gru2_b')(gru1_merged)\n",
    "x = concatenate([gru_2, gru_2b])\n",
    "x = Dropout(0.25)(x)\n",
    "x = Dense(n_class, activation='softmax', kernel_regularizer=l2(l2_rate))(x)\n",
    "rnn_out = x\n",
    "base_model = Model(input_tensor, x)\n",
    "\n",
    "base_model2 = make_parallel(base_model, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 预测\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def disp3(index):\n",
    "    s = decode(out[index])\n",
    "    plt.figure(figsize=(16, 4))\n",
    "    plt.imshow(X[index].transpose(1, 0, 2))\n",
    "    plt.title('pred:%s'%s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = '0.997656'\n",
    "base_model.load_weights('model_l2加层生成器2_%s.h5' % z)\n",
    "y_pred = base_model2.predict(X, batch_size=500, verbose=1)\n",
    "out = K.get_value(K.ctc_decode(y_pred[:,2:], input_length=np.ones(y_pred.shape[0])*rnn_length)[0][0])[:, :n_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ss = map(decode, out)\n",
    "\n",
    "vals = []\n",
    "errs = []\n",
    "for i in tqdm(range(100000)):\n",
    "    a = ss[i].split(';')\n",
    "    s = a[-1]\n",
    "    for x in a[:-1]:\n",
    "        x, c = x.split('=')\n",
    "        s = s.replace(x, c+'.0')\n",
    "    \n",
    "    val = ''\n",
    "    try:\n",
    "        val = '%.2f' % eval(s)\n",
    "    except:\n",
    "        ss[i] = ''\n",
    "        errs.append(s)\n",
    "    \n",
    "    vals.append(val)\n",
    "    \n",
    "with open('result_%s.txt' % z, 'w') as f:\n",
    "    f.write('\\n'.join(map(' '.join, list(zip(ss, vals)))).encode('utf-8'))\n",
    "    \n",
    "print len(errs)\n",
    "print 1-len(errs)/100000."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
